# -*- coding: utf-8 -*-
"""
Created on Tue Dec 21 11:59:59 2021

@author: helai.bai
"""
import json
import pandas
import xgboost

class xgb_parser(object):
    '''
    This method is used to parse the xgboost.sklearn.XGBClassifier model into if-else decision trees.
    
    To use this method, firstly, you should have a model, or load from your saved file.
        
        import joblib
        model = joblib.load('xgb_msxj_1125.pkl')
    
    Note: Your model must have feature names stored. You can check it with this:
        
        model.get_booster().feature_names
        model.get_booster().feature_types
        
    If not, you could manually set your model features by this:
        
        features = [feature_1, feature_2]
        model.get_booster().feature_names = features
        model.get_booster().feature_types = ['float' for _ in features]
        
    Then, use this class to parse your model, and save the generated code into your given path.
        
        from xgb_parser import xgb_parser
        
        parser = xgb_parser(model = model)
        parser.generate_if_else_trees('./xgb_model.py')
        
    To use the generated code, you should import the 'predict' method from the generated model file. 
        
        from xgb_model import predict
        
    The 'predict' method takes **one** sample and output its probability.
    Your data must be in **dictionary** type and only float numbers are accepted. 
    Be aware that, 'nan' is not supported, you should **convert it into 'None'** as follows.
        
        xgb_data_df = df[features_in_order].iloc[1:2]
        
        value_dict = {}
        for feature in f:
            value = xgb_data_df[feature].iloc[0]
            if numpy.isnan(value):
                value_dict[feature] = None
            else:
                value_dict[feature] = float(value)
                
    Or, you can use the provided static method in this class to achieve that.
        
        parser.data_wrapper(data[features])

    Then use the 'predict' method.
    
        res = []
        for value_dict in parser.data_wrapper(data[features]):
            proba = predict(value_dict)
            res.append(proba)
    '''
    
    def __init__(self, model):
        assert isinstance(model, xgboost.sklearn.XGBClassifier), 'The provided model has unsupported type.'
        
        self.model = model
        self.features = model.get_booster().feature_names
        assert self.features is not None
        
    def _get_dump_strings(self) -> [str]:
        '''
        Get the decision tree json version.

        Returns
        -------
        [str]
            List of json strings of sub decision trees.

        '''
        return self.model.get_booster().get_dump(dump_format='json')
    
    def _generate_node_string(self, tree : dict) -> ([str], int):
        '''
        Convert the current tree node into if-else block.

        Parameters
        ----------
        tree : dict
            DESCRIPTION.

        Returns
        -------
        [str]
            The lines of sub if-else block in the current node.
        int
            The node id

        '''
    
        # ==============
        # leaf node
        # ==============
        if 'leaf' in tree:
            return ["return " + str(tree['leaf'])], tree['nodeid']
            
        # ==============
        # non-leaf node
        # ==============
        
        # curr feature 
        feature = tree['split']
        
        # get left and right sub tree
        left_tree, right_tree = tree['children']
        left_str,  left_id  = self._generate_node_string(left_tree)
        right_str, right_id = self._generate_node_string(right_tree)
        
        # add indentation of the sub tree
        left_str  = ['    ' + line for line in left_str]
        right_str = ['    ' + line for line in right_str]
        
        # generate the current if-else
        if tree['yes'] == tree['missing']:
            if_clause = "if value_dict['" + feature + "'] is None or value_dict['" + feature + "'] < " + str(tree['split_condition']) + ":"
        else:
            if_clause = "if value_dict['" + feature + "'] is None or value_dict['" + feature + "'] >= " + str(tree['split_condition']) + ":"
            left_str, right_str = right_str, left_str
        else_clause = 'else:'
        
        curr_tree_str = [if_clause] + left_str + [else_clause] + right_str 
        
        return curr_tree_str, tree['nodeid']
    
    def generate_if_else_trees(self, path : str) -> None:
        '''
        Generate the if-else version of the current xgboost model.
        Each node in the decison trees is converted into a nested if-else block.

        Parameters
        ----------
        path : str
            The file path to output the generated code.

        Returns
        -------
        None
            DESCRIPTION.

        '''
        
        tree_strings = self._get_dump_strings()
        
        final_strings = []
        
        # sub trees
        count = 0
        for tree_string in tree_strings:
            tree = json.loads(tree_string)
            tree_str = self._generate_node_string(tree)[0]
            tree_str = ['    ' + line for line in tree_str]
            tree_str = ['def sub_decision_tree_' + str(count) + '(value_dict):'] + tree_str + ['']
            count += 1
            
            final_strings += tree_str
        
        # summation
        base_point = 0
        summation_line = 'score = ' + str(base_point)
        for i in range(0, count):
            summation_line = summation_line + ' + sub_decision_tree_' + str(i) + '(value_dict)'
            
        final_strings = final_strings + \
                        ['def predict(value_dict):'] + \
                        ['    ' + summation_line] + \
                        ['    from math import exp'] + \
                        ['    return 1.0 / (1.0 + exp(-score))']
            
        final_strings = [line + '\n' for line in final_strings]
        
        with open(path, 'w') as file_writer:
            file_writer.writelines(final_strings)
            
    
    @staticmethod
    def data_wrapper(data : pandas.core.frame.DataFrame) -> dict:
        '''
        Convert a dataframe into a dictionary type which is valid for the class.
        Convert all numeric value into float type and nan into None.
        Each line in the dataframe is converted into a dictionary.
        This function is used as an iterator.

        Parameters
        ----------
        data : pandas.core.frame.DataFrame
            The input dataframe.

        Returns
        -------
        dict
            The dictionary value for 'predict' method.

        '''
        if isinstance(data, pandas.core.frame.DataFrame):
            for i in range(data.shape[0]):
                value_dict = data.iloc[i].to_dict()
                
                for feature in value_dict.keys():
                    value = value_dict[feature]
                    if value is None or numpy.isnan(value):
                        value_dict[feature] = None
                    else:
                        value_dict[feature] = float(value)                        
                
                yield value_dict